#!/bin/env python
# -*- coding:utf-8
"""
@author:xingz
@file:queens.py
@time:2022/05/03
"""
from typing import List, Dict
import matplotlib.pyplot as plt
from zx_constraint.csp import Constraint, CSP


class QueenConstraint(Constraint[int, int]):
    def __init__(self, columns: List[int]):
        super(QueenConstraint, self).__init__(columns)
        self.columns = columns

    def satisfied(self, assignment: Dict[int, int]) -> bool:
        for c, r in assignment.items():
            for c2, r2 in {k: v for k, v in assignment.items() if k != c}.items():
                if r == r2:
                    return False
                if abs(c2 - c) == abs(r2 - r):
                    return False
        return True


variables: List[int] = [1, 2, 3, 4, 5, 6, 7, 8]
domains: Dict[int, List[int]] = {}
for v in variables:
    domains[v] = [1, 2, 3, 4, 5, 6, 7, 8]

csp = CSP(variables=variables, domains=domains)
csp.add_constraint(QueenConstraint(variables))
solution = csp.backtracking_search()
print(solution)
plt.scatter(solution.keys(), solution.values())
plt.grid(b=True)
plt.show()